#!/usr/bin/env python3
# -*- coding:utf-8 -*-
"""
CIFAR10 图像分类预测脚本
支持加载预训练模型并对单张图片进行预测
"""

import argparse
from pathlib import Path

import torch
import torchvision.transforms as transforms
from PIL import Image

from config import config
from models.cnn import CIFAR10CNN
from utils.logger import setup_logger

current_file = Path(__file__).name
logger = setup_logger(current_file)

# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_model(model_path: str) -> CIFAR10CNN:
    """
    加载预训练模型
    Args:
        model_path: 模型文件路径（如 ./models/model-cifar10.pth）
    Returns:
        加载后的模型实例
    """
    # 初始化模型
    model = CIFAR10CNN(num_classes=10).to(config.DEVICE)

    # 加载训练好的模型
    try:
        checkpoint = torch.load(model_path)
        model.load_state_dict(checkpoint['model_state_dict'])
        model.eval()
        logger.info("Model loaded successfully")
    except FileNotFoundError:
        logger.error(f"No trained model found at {config.BEST_MODEL_PATH}")
    return model


def preprocess_image(img_path: str) -> torch.Tensor:
    """
    图像预处理（与训练时保持一致）
    Args:
        img_path: 输入图像路径
    Returns:
        归一化后的4D张量（1,3,32,32）
    """
    transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
    ])
    img = Image.open(img_path).convert('RGB')
    img = transform(img).unsqueeze(0)  # 增加batch维度
    return img.to(device)


def predict(model: torch.nn.Module, img_tensor: torch.Tensor) -> tuple:
    """
    执行预测
    Args:
        model: 预训练模型
        img_tensor: 预处理后的图像张量
    Returns:
        (预测类别索引, 预测概率)
    """
    with torch.no_grad():
        output = model(img_tensor)
        prob = torch.softmax(output, dim=1)
        pred_idx = torch.argmax(prob).item()
        confidence = prob[0, pred_idx].item()
        return pred_idx, confidence


def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="CIFAR10 图像分类预测")
    parser.add_argument("--model", default=config.BEST_MODEL_PATH, type=str, required=False, help="模型文件路径")
    parser.add_argument("--image", type=str, required=True, help="待预测图像路径")
    args = parser.parse_args()

    # 加载模型和图像
    try:
        model = load_model(args.model)
        img_tensor = preprocess_image(args.image)
    except Exception as e:
        print(f"加载失败: {str(e)}")
        return

    # 执行预测
    pred_idx, confidence = predict(model, img_tensor)

    # 结果可视化
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    print(f"预测类别: {classes[pred_idx]}\n置信度: {confidence:.2%}")


if __name__ == "__main__":
    main()
